package org.springframework.jdbc.datasource;

import java.sql.Connection;

import javax.sql.DataSource;

import org.springframework.transaction.CannotCreateTransactionException;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.support.DefaultTransactionStatus;
import org.springframework.transaction.support.TransactionSynchronizationManager;

import com.springboot.distributed.transaction.DistributedTransactionDataSourceAdapter;

/**
 * 分布式数据源事务管理器
 * <p>
 * 注意：此包名必须和spring源码中的一致，不然 ConnectionHolder.transactionActive 属性无法设置
 *
 * @author supu
 * @date 2019-09-06 15:34
 **/
public class DistributedDataSourceTransactionManager extends DataSourceTransactionManager {
    private static final long serialVersionUID = 1043411192408908855L;

    private DistributedTransactionDataSourceAdapter distributedTransactionDataSourceAdapter;

    public DistributedDataSourceTransactionManager() {
        super();
    }

    public DistributedDataSourceTransactionManager(DataSource dataSource) {
        super(dataSource);
    }

    public DistributedTransactionDataSourceAdapter getDistributedTransactionDataSourceAdapter() {
        return distributedTransactionDataSourceAdapter;
    }

    public void setDistributedTransactionDataSourceAdapter(DistributedTransactionDataSourceAdapter distributedTransactionDataSourceAdapter) {
        this.distributedTransactionDataSourceAdapter = distributedTransactionDataSourceAdapter;
    }

    @Override
    protected Object doGetTransaction() {
        DistributedDataSourceTransactionObject txObject = new DistributedDataSourceTransactionObject();
        txObject.setSavepointAllowed(isNestedTransactionAllowed());
        ConnectionHolder conHolder =
                (ConnectionHolder) TransactionSynchronizationManager.getResource(obtainDataSource());
        txObject.setConnectionHolder(conHolder, false);
        return txObject;
    }

    @Override
    protected boolean isExistingTransaction(Object transaction) {
        DistributedDataSourceTransactionObject txObject = (DistributedDataSourceTransactionObject) transaction;
        return (txObject.hasConnectionHolder() && txObject.getConnectionHolder().isTransactionActive());
    }

    /**
     * This implementation sets the isolation level but ignores the timeout.
     */
    @Override
    protected void doBegin(Object transaction, TransactionDefinition definition) {
        DistributedDataSourceTransactionObject txObject = (DistributedDataSourceTransactionObject) transaction;
        Connection con = null;

        try {
            if (!txObject.hasConnectionHolder() ||
                    txObject.getConnectionHolder().isSynchronizedWithTransaction()) {
                Connection newCon = obtainDataSource().getConnection();
                if (logger.isDebugEnabled()) {
                    logger.debug("Acquired Connection [" + newCon + "] for JDBC transaction");
                }
                txObject.setConnectionHolder(new ConnectionHolder(newCon), true);
            }

            txObject.getConnectionHolder().setSynchronizedWithTransaction(true);
            con = txObject.getConnectionHolder().getConnection();

            Integer previousIsolationLevel = DataSourceUtils.prepareConnectionForTransaction(con, definition);
            txObject.setPreviousIsolationLevel(previousIsolationLevel);

            // Switch to manual commit if necessary. This is very expensive in some JDBC drivers,
            // so we don't want to do it unnecessarily (for example if we've explicitly
            // configured the connection pool to set it already).
            if (con.getAutoCommit()) {
                txObject.setMustRestoreAutoCommit(true);
                if (logger.isDebugEnabled()) {
                    logger.debug("Switching JDBC Connection [" + con + "] to manual commit");
                }
                con.setAutoCommit(false);
            }

            prepareTransactionalConnection(con, definition);
            txObject.getConnectionHolder().setTransactionActive(true);

            int timeout = determineTimeout(definition);
            if (timeout != TransactionDefinition.TIMEOUT_DEFAULT) {
                txObject.getConnectionHolder().setTimeoutInSeconds(timeout);
            }

            // Bind the connection holder to the thread.
            if (txObject.isNewConnectionHolder()) {
                TransactionSynchronizationManager.bindResource(obtainDataSource(), txObject.getConnectionHolder());
            }
        } catch (Throwable ex) {
            if (txObject.isNewConnectionHolder()) {
                DataSourceUtils.releaseConnection(con, obtainDataSource());
                txObject.setConnectionHolder(null, false);
            }
            throw new CannotCreateTransactionException("Could not open JDBC Connection for transaction", ex);
        }
    }

    @Override
    protected Object doSuspend(Object transaction) {
        DistributedDataSourceTransactionObject txObject = (DistributedDataSourceTransactionObject) transaction;
        txObject.setConnectionHolder(null);
        return TransactionSynchronizationManager.unbindResource(obtainDataSource());
    }

    @Override
    protected void doCommit(DefaultTransactionStatus status) {
        distributedTransactionDataSourceAdapter.doCommit(status);
    }

    @Override
    protected void doRollback(DefaultTransactionStatus status) {
        distributedTransactionDataSourceAdapter.doRollBack(status);
    }

    @Override
    protected void doSetRollbackOnly(DefaultTransactionStatus status) {
        DistributedDataSourceTransactionObject txObject = (DistributedDataSourceTransactionObject) status.getTransaction();
        if (status.isDebug()) {
            logger.debug("Setting JDBC transaction [" + txObject.getConnectionHolder().getConnection() +
                    "] rollback-only");
        }
        txObject.setRollbackOnly();
    }

    @Override
    protected void doCleanupAfterCompletion(Object transaction) {
        DistributedDataSourceTransactionObject txObject = (DistributedDataSourceTransactionObject) transaction;

        // Remove the connection holder from the thread, if exposed.
        if (txObject.isNewConnectionHolder()) {
            TransactionSynchronizationManager.unbindResource(obtainDataSource());
        }

        distributedTransactionDataSourceAdapter.doCleanupAfterCompletion(txObject);
    }
}
